﻿using System;
using System.Buffers;
using System.Collections.Generic;
using System.Diagnostics;
using System.IO;
using System.Runtime.CompilerServices;
using System.Runtime.InteropServices;
using System.Threading;
using System.Threading.Tasks;
using PooledGrowableBufferHelper;
using SixLabors.ImageSharp;
using SixLabors.ImageSharp.PixelFormats;
using WholeSlideImageLibrary.DeepZoom;
using WholeSlideImageLibrary.Features;

namespace WholeSlideImageLibrary
{
    public sealed class WholeSlideImageShim : WholeSlideImage,
        ITextPropertiesFeature, IDeepZoomProfileHintFeature, IDeepZoomFeature, IAssociatedImagesFeature, IAssociatedImagesEncodingFeature
    {
        private readonly WholeSlideImage _image;

        private WholeSlideImageShim(WholeSlideImage image)
        {
            _image = image ?? throw new ArgumentNullException(nameof(image));
        }


        public static WholeSlideImageShim Wrap(WholeSlideImage image)
        {
            if (image is null)
            {
                throw new ArgumentNullException(nameof(image));
            }
            if (image is WholeSlideImageShim shim)
            {
                return shim;
            }
            return new WholeSlideImageShim(image);
        }
        public override int LevelCount => _image.LevelCount;

        public override int Width => _image.Width;

        public override int Height => _image.Height;

        public override double MicronsPerPixelX => _image.MicronsPerPixelX;

        public override double MicronsPerPixelY => _image.MicronsPerPixelY;

        public override System.Drawing.Size GetLevelSize(int level)
            => _image.GetLevelSize(level);

        public override ValueTask ReadRegionAsync<T>(int level, int x, int y, int width, int height, Memory<T> buffer, CancellationToken cancellationToken)
            => _image.ReadRegionAsync(level, x, y, width, height, buffer, cancellationToken);

        #region ITextPropertiesFeature

        public ValueTask<IEnumerable<string>> GetAllTextPropertiesAsync(CancellationToken cancellationToken = default)
        {
            if (_image is ITextPropertiesFeature feature)
            {
                return feature.GetAllTextPropertiesAsync(cancellationToken);
            }

            return new ValueTask<IEnumerable<string>>(Array.Empty<string>());
        }

        public ValueTask<string> ReadTextPropertyAsync(string name, CancellationToken cancellationToken = default)
        {
            if (_image is ITextPropertiesFeature feature)
            {
                return feature.ReadTextPropertyAsync(name, cancellationToken);
            }

            return ThrowKeyNotFoundException<string>();
        }

        #endregion


        #region IDeepZoomProfileHintFeature and IDeepZoomFeature

        public IEnumerable<DeepZoomProfile> GetDeepZoomProfileHints()
        {
            if (_image is IDeepZoomProfileHintFeature feature)
            {
                return feature.GetDeepZoomProfileHints();
            }

            return Array.Empty<DeepZoomProfile>();
        }

        public async ValueTask<DeepZoomTileSource> OpenDeepZoomAsync(DeepZoomProfile profile, CancellationToken cancellationToken = default)
        {
            DeepZoomTileSource? deepZoom = null;
            if (_image is IDeepZoomFeature feature)
            {
                deepZoom = await feature.OpenDeepZoomAsync(profile, cancellationToken).ConfigureAwait(false);
            }
            if (deepZoom is null)
            {
                deepZoom = new DefaultDeepZoomTileSource(_image, profile);
            }

            return deepZoom;
        }

        async ValueTask<DeepZoomTileSource?> IDeepZoomFeature.OpenDeepZoomAsync(DeepZoomProfile profile, CancellationToken cancellationToken)
        {
            DeepZoomTileSource? deepZoom = null;
            if (_image is IDeepZoomFeature feature)
            {
                deepZoom = await feature.OpenDeepZoomAsync(profile, cancellationToken).ConfigureAwait(false);
            }
            if (deepZoom is null)
            {
                deepZoom = new DefaultDeepZoomTileSource(_image, profile);
            }

            return deepZoom;
        }

        #endregion

        #region Associated Images
        public ValueTask<IEnumerable<string>> GetAllAssociatedImagesAsync(CancellationToken cancellationToken = default)
        {
            if (_image is IAssociatedImagesFeature feature)
            {
                return feature.GetAllAssociatedImagesAsync(cancellationToken);
            }
            if (_image is IAssociatedImagesEncodingFeature encodingFeature)
            {
                return encodingFeature.GetAllAssociatedImagesAsync(cancellationToken);
            }

            return new ValueTask<IEnumerable<string>>(Array.Empty<string>());
        }

        public ValueTask<System.Drawing.Size> GetAssociatedImageSizeAsync(string name, CancellationToken cancellationToken = default)
        {
            if (_image is IAssociatedImagesFeature feature)
            {
                return feature.GetAssociatedImageSizeAsync(name, cancellationToken);
            }
            if (_image is IAssociatedImagesEncodingFeature encodingFeature)
            {
                return encodingFeature.GetAssociatedImageSizeAsync(name, cancellationToken);
            }

            return ThrowKeyNotFoundException<System.Drawing.Size>();
        }

        public ValueTask ReadAssociatedImageAsync<T>(string name, Memory<T> buffer, CancellationToken cancellationToken = default) where T : unmanaged
        {
            if (_image is IAssociatedImagesFeature feature)
            {
                return feature.ReadAssociatedImageAsync(name, buffer, cancellationToken);
            }

            if (_image is IAssociatedImagesEncodingFeature encodingFeature)
            {
                return new ValueTask(ReadAssociatedImageAsyncCore(encodingFeature, name, buffer, cancellationToken));
            }

            return ThrowKeyNotFoundException();

            static async Task ReadAssociatedImageAsyncCore(IAssociatedImagesEncodingFeature encodingFeature, string name, Memory<T> buffer, CancellationToken cancellationToken)
            {
                System.Drawing.Size size = await encodingFeature.GetAssociatedImageSizeAsync(name, cancellationToken).ConfigureAwait(false);
                EncodedImageFileFormat fileFormat = await encodingFeature.GetAssociatedImageEncodingSupportAsync(name, cancellationToken).ConfigureAwait(false);

                using (Stream stream = PooledMemoryStreamManager.Shared.GetStream())
                {
                    if ((fileFormat & EncodedImageFileFormat.Jpeg) != 0)
                    {
                        fileFormat = EncodedImageFileFormat.Jpeg;
                        await encodingFeature.EncodeAssociatedImageAsync(name, fileFormat, stream, cancellationToken).ConfigureAwait(false);
                    }
                    else if ((fileFormat & EncodedImageFileFormat.Png) != 0)
                    {
                        fileFormat = EncodedImageFileFormat.Png;
                        await encodingFeature.EncodeAssociatedImageAsync(name, fileFormat, stream, cancellationToken).ConfigureAwait(false);
                    }
                    else if ((fileFormat & EncodedImageFileFormat.Bmp) != 0)
                    {
                        fileFormat = EncodedImageFileFormat.Bmp;
                        await encodingFeature.EncodeAssociatedImageAsync(name, fileFormat, stream, cancellationToken).ConfigureAwait(false);
                    }
                    else
                    {
                        throw new NotSupportedException();
                    }

                    stream.Seek(0, SeekOrigin.Begin);
                    using var image = Image.Load<Rgba32>(stream);
                    if (Unsafe.SizeOf<T>() * buffer.Length < 4 * image.Width * image.Height)
                    {
                        throw new ArgumentException("destination is too small.", nameof(buffer));
                    }

                    CopyPixels(image, buffer);
                }
            }
        }

        public async ValueTask<EncodedImageFileFormat> GetAssociatedImageEncodingSupportAsync(string name, CancellationToken cancellationToken)
        {
            EncodedImageFileFormat format = EncodedImageFileFormat.None;
            if (_image is IAssociatedImagesEncodingFeature encodingFeature)
            {
                format = await encodingFeature.GetAssociatedImageEncodingSupportAsync(name, cancellationToken).ConfigureAwait(false);
            }

            if (_image is IAssociatedImagesFeature && format == EncodedImageFileFormat.None)
            {
                format = EncodedImageFileFormat.Jpeg;
            }

            return format;
        }

        public ValueTask EncodeAssociatedImageAsync(string name, EncodedImageFileFormat format, Stream stream, CancellationToken cancellationToken = default)
        {
            if (_image is IAssociatedImagesEncodingFeature encodingFeature)
            {
                return new ValueTask(EncodeAssociatedImageAsyncCore(encodingFeature, name, format, stream, cancellationToken));
            }

            if (_image is IAssociatedImagesFeature feature)
            {
                return new ValueTask(EncodeFromBufferAsyncCore(feature, name, format, stream, cancellationToken));
            }

            return ThrowKeyNotFoundException();

            static async Task EncodeAssociatedImageAsyncCore(IAssociatedImagesEncodingFeature encodingFeature, string name, EncodedImageFileFormat format, Stream stream, CancellationToken cancellationToken)
            {
                EncodedImageFileFormat supportedFeature = await encodingFeature.GetAssociatedImageEncodingSupportAsync(name, cancellationToken).ConfigureAwait(false);
                if ((supportedFeature & format) != 0)
                {
                    await encodingFeature.EncodeAssociatedImageAsync(name, format, stream, cancellationToken).ConfigureAwait(false);
                }
                else if (encodingFeature is IAssociatedImagesFeature feature)
                {
                    await EncodeFromBufferAsyncCore(feature, name, format, stream, cancellationToken).ConfigureAwait(false);
                }
                else
                {
                    throw new NotSupportedException();
                }
            }

            static async Task EncodeFromBufferAsyncCore(IAssociatedImagesFeature feature, string name, EncodedImageFileFormat format, Stream stream, CancellationToken cancellationToken = default)
            {
                System.Drawing.Size size = await feature.GetAssociatedImageSizeAsync(name, cancellationToken).ConfigureAwait(false);
                byte[] buffer = ArrayPool<byte>.Shared.Rent(4 * size.Width * size.Height);
                try
                {
                    Memory<Rgba32> memory = buffer.AsMemory(0, 4 * size.Width * size.Height).Cast<byte, Rgba32>();
                    using var image = Image.WrapMemory(memory, size.Width, size.Height);
                    switch (format)
                    {
                        case EncodedImageFileFormat.Bmp:
                            image.SaveAsBmp(stream);
                            break;
                        case EncodedImageFileFormat.Jpeg:
                            image.SaveAsJpeg(stream);
                            break;
                        case EncodedImageFileFormat.Png:
                            image.SaveAsPng(stream);
                            break;
                        default:
                            throw new NotSupportedException();
                    }
                }
                finally
                {
                    ArrayPool<byte>.Shared.Return(buffer);
                }
            }
        }

        public ValueTask EncodeAssociatedImageAsync(string name, EncodedImageFileFormat format, IBufferWriter<byte> writer, CancellationToken cancellationToken)
        {
            if (_image is IAssociatedImagesEncodingFeature encodingFeature)
            {
                return new ValueTask(EncodeAssociatedImageAsyncCore(encodingFeature, name, format, writer, cancellationToken));
            }

            if (_image is IAssociatedImagesFeature feature)
            {
                return new ValueTask(EncodeFromBufferAsyncCore(feature, name, format, writer, cancellationToken));
            }

            return ThrowKeyNotFoundException();

            static async Task EncodeAssociatedImageAsyncCore(IAssociatedImagesEncodingFeature encodingFeature, string name, EncodedImageFileFormat format, IBufferWriter<byte> writer, CancellationToken cancellationToken)
            {
                EncodedImageFileFormat supportedFeature = await encodingFeature.GetAssociatedImageEncodingSupportAsync(name, cancellationToken).ConfigureAwait(false);
                if ((supportedFeature & format) != 0)
                {
                    await encodingFeature.EncodeAssociatedImageAsync(name, format, writer, cancellationToken).ConfigureAwait(false);
                }
                else if (encodingFeature is IAssociatedImagesFeature feature)
                {
                    await EncodeFromBufferAsyncCore(feature, name, format, writer, cancellationToken).ConfigureAwait(false);
                }
                else
                {
                    throw new NotSupportedException();
                }
            }

            static async Task EncodeFromBufferAsyncCore(IAssociatedImagesFeature feature, string name, EncodedImageFileFormat format, IBufferWriter<byte> writer, CancellationToken cancellationToken)
            {
                System.Drawing.Size size = await feature.GetAssociatedImageSizeAsync(name, cancellationToken).ConfigureAwait(false);
                using PooledMemoryStream? stream = PooledMemoryStreamManager.Shared.GetStream();
                byte[] buffer = ArrayPool<byte>.Shared.Rent(4 * size.Width * size.Height);
                try
                {
                    Memory<Rgba32> memory = buffer.AsMemory(0, 4 * size.Width * size.Height).Cast<byte, Rgba32>();
                    using var image = Image.WrapMemory(memory, size.Width, size.Height);
                    switch (format)
                    {
                        case EncodedImageFileFormat.Bmp:
                            image.SaveAsBmp(stream);
                            break;
                        case EncodedImageFileFormat.Jpeg:
                            image.SaveAsJpeg(stream);
                            break;
                        case EncodedImageFileFormat.Png:
                            image.SaveAsPng(stream);
                            break;
                        default:
                            throw new NotSupportedException();
                    }
                }
                finally
                {
                    ArrayPool<byte>.Shared.Return(buffer);
                }

                stream.Seek(0, SeekOrigin.Begin);
                Copy(stream.ToReadOnlySequence(), writer);
            }
        }

        #endregion

        private static ValueTask<T> ThrowKeyNotFoundException<T>()
        {
            throw new KeyNotFoundException();
        }

        private static ValueTask ThrowKeyNotFoundException()
        {
            throw new KeyNotFoundException();
        }

        private static void CopyPixels<T>(Image<Rgba32> source, Memory<T> destination) where T : unmanaged
        {
            Debug.Assert(Unsafe.SizeOf<T>() * destination.Length >= 4 * source.Width * source.Height);

            Span<byte> destinationSpan = MemoryMarshal.AsBytes(destination.Span);

            if (source.TryGetSinglePixelSpan(out Span<Rgba32> sourceSpan))
            {
                MemoryMarshal.AsBytes(sourceSpan).CopyTo(destinationSpan.Slice(0, 4 * sourceSpan.Length));
            }
            else
            {
                int height = source.Height;
                int rowByteCount = 4 * source.Width;
                for (int i = 0; i < height; i++)
                {
                    Span<Rgba32> rowSpan = source.GetPixelRowSpan(i);
                    MemoryMarshal.AsBytes(rowSpan).CopyTo(destinationSpan.Slice(0, rowByteCount));
                    destinationSpan = destinationSpan.Slice(rowByteCount);
                }
            }
        }

        const int MaxSegmentSize = 16384;

        private static void Copy(ReadOnlySequence<byte> sequence, IBufferWriter<byte> writer)
        {
            long length = sequence.Length;

            while (length != 0)
            {
                int segmentSize = (int)Math.Min(length, MaxSegmentSize);
                Span<byte> span = writer.GetSpan(segmentSize);
                segmentSize = (int)Math.Min(span.Length, length);
                sequence.Slice(0, segmentSize).CopyTo(span);
                writer.Advance(segmentSize);
                length -= segmentSize;
            }
        }

    }
}
